# preprocess.py
import pandas as pd
import joblib
import os
from sklearn.preprocessing import StandardScaler

# 数值特征列表（已移除id）
NUMERIC_COLUMNS = ["预算", "流行度", "平均评分", "评分人数", "已上映天数", "入座率", "影院评分"]


def preprocess_and_save(train_dir, test_dir, save_dir):
    os.makedirs(save_dir, exist_ok=True)

    def load_all_csv(directory):
        df_list = []
        for f in sorted(os.listdir(directory)):
            if f.endswith(".csv"):
                file_path = os.path.join(directory, f)
                try:
                    df = pd.read_csv(file_path, encoding='utf-8')
                except UnicodeDecodeError:
                    df = pd.read_csv(file_path, encoding='gbk')
                df_list.append(df)
        return pd.concat(df_list, ignore_index=True)

    print("📥 加载数据...")
    train_df = load_all_csv(train_dir)
    test_df = load_all_csv(test_dir)

    # 标准化数值特征
    print("⚙️  标准化特征...")
    scaler = StandardScaler()
    train_features = scaler.fit_transform(train_df[NUMERIC_COLUMNS])
    test_features = scaler.transform(test_df[NUMERIC_COLUMNS])

    joblib.dump(scaler, os.path.join(save_dir, "scaler.pkl"))

    # 保存为pkl
    print("💾 保存为pkl...")
    joblib.dump({"features": train_features, "labels": train_df["历史票价"].values}, os.path.join(save_dir, "train.pkl"))
    joblib.dump({"features": test_features, "labels": test_df["历史票价"].values}, os.path.join(save_dir, "test.pkl"))

    print("✅ 数据预处理完成")


if __name__ == "__main__":
    train_path = "jingjia/data/train"
    test_path = "jingjia/data/test"
    save_path = "jingjia/data/processed"

    preprocess_and_save(train_path, test_path, save_path)

